#include "Rdma.h"

#include <string.h>

#include <iostream>

using std::cout;
using std::cin;
using std::endl;

struct ibv_context *Rdma::dev_ctx_ = nullptr;
struct ibv_pd *Rdma::pd_ = nullptr;
static int ref = 0;
Rdma::Rdma(MemoryManager &mm) : mm_(mm)
{}
Rdma::~Rdma()
{}
void Rdma::Destroy()
{
    ref--;
    if(srq_)
    {
        ibv_destroy_srq(srq_);
    }
    if(!ref && pd_)
    {
        ibv_dealloc_pd(pd_);
    }
    if(send_cq_)
    {
        ibv_destroy_cq(send_cq_);
    }
    if(recv_cq_)
    {
        ibv_destroy_cq(recv_cq_);
        
    }
    
    if(!ref && dev_ctx_)
    {   
        ibv_close_device(dev_ctx_);
    }
    
}

int Rdma::Init(const char* dev_name, uint8_t port_id, uint8_t flag)
{
    int dev_num = 0;
    struct ibv_device** dev_list = NULL;
    struct ibv_device* dev = NULL;


    bool is_resolved = false;
    uint8_t resolved_port;
    /* get device list*/
    dev_list = ibv_get_device_list(&dev_num);
    if (!dev_list) {
        cout << "Failed to get device list!!!" << endl;
        return -1;
    }
    if (!dev_num) {
        cout << "Get " << dev_num << " devices" << endl;
        return -1;
    }

    /* resolve device and port*/
    for (int dev_i = 0; dev_i < dev_num; dev_i++) {
        /* reslove device*/
        if (!dev_name) {
            dev = dev_list[dev_i];
        }
        else {
            if (!strcmp(ibv_get_device_name(dev_list[dev_i]), dev_name)) {
                dev = dev_list[dev_i];
            }
        }

        /* resolve port*/
        if (dev) {
            struct ibv_context *dev_ctx = ibv_open_device(dev);
            if (!dev_ctx) {
                cout << "Failed to open device" << endl;
                break;
            }
            
            memset(&dev_attr_, 0, sizeof(dev_attr_));
            if (ibv_query_device(dev_ctx, &dev_attr_)) {
                cout << "Failed to query device" << endl;
                break;
            }

            memset(&port_attr_, 0, sizeof(port_attr_));
            if (port_id > 0) {
                if (ibv_query_port(dev_ctx, port_id, &port_attr_)) {
                    cout << "Failed to query port" << endl;
                    break;
                }
                if (port_attr_.state == IBV_PORT_ACTIVE) {
                    resolved_port = port_id;
                    is_resolved = true;
                    break;
                }
            }

            for (uint8_t port_i = 1; port_i <= dev_attr_.phys_port_cnt; port_i++) {
                memset(&port_attr_, 0, sizeof(port_attr_));
                if (ibv_query_port(dev_ctx, port_i, &port_attr_)) {
                    cout << "Failed to query port" << endl;
                    break;
                }
                if (port_attr_.state == IBV_PORT_ACTIVE) {
                    resolved_port = port_i;
                    is_resolved = true;
                    break;
                }
            }

            if (is_resolved) {
                if(!dev_ctx_)
                {
                    ref++;
                    dev_ctx_ = dev_ctx;
                }
                break;
            }
        }
    }

    /* dev_list is useless, free it*/
    ibv_free_device_list(dev_list);

    if (!is_resolved) {
        cout << "Failed to resolve device and port" << endl;
        ibv_close_device(dev_ctx_);
        return -1;
    }
    resolve_port_id_ = resolved_port;

    /* create completion queue */
    send_cq_ = ibv_create_cq(dev_ctx_, CQ_MAX_DEPTH, NULL, NULL, 0);
    recv_cq_ = ibv_create_cq(dev_ctx_, CQ_MAX_DEPTH, NULL, NULL, 0);
    if (!send_cq_ || !recv_cq_) {
        cout << "Failed to create cq" << endl;
        ibv_close_device(dev_ctx_);
        return -1;
    }

    /* create protection domain */
    if(!pd_)
    {
        pd_ = ibv_alloc_pd(dev_ctx_);

        if (!pd_) {
            cout << "Failed to malloc pd" << endl;
            ibv_destroy_cq(send_cq_);
            ibv_destroy_cq(recv_cq_);
            ibv_close_device(dev_ctx_);
            return -1;
        }
    }
    /* create srq*/
    if(ENABLE_SRQ & flag)
    {
        struct ibv_srq_init_attr srq_init_attr;

        memset(&srq_init_attr, 0, sizeof(srq_init_attr));

        srq_init_attr.attr.max_wr = QP_MAX_DEPTH;
        srq_init_attr.attr.max_sge = WR_MAX_SGE;

        srq_ = ibv_create_srq(pd_, &srq_init_attr);
        if (!srq_) {
            ibv_dealloc_pd(pd_);
            ibv_destroy_cq(send_cq_);
            ibv_destroy_cq(recv_cq_);
            ibv_close_device(dev_ctx_);
            cout << "failed to create shard recv queue" << endl;
            return -1;
        }
    }
    /* register memeory */
    if(!mm_.mr())
    {
        uint64_t memory_size = mm_.memory_size();
        char *memory = (char *)mm_.memory();
        int mr_flag = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_ATOMIC;
        cout << "pd : " << pd_ << "memory : " << memory << "memory_size : " << memory_size << endl; 
        struct ibv_mr *mr = ibv_reg_mr(pd_, memory, memory_size, mr_flag);
        
        if (!mr) {
            exit(-1);
            ibv_destroy_srq(srq_);
            ibv_dealloc_pd(pd_);
            ibv_destroy_cq(send_cq_);
            ibv_destroy_cq(recv_cq_);
            ibv_close_device(dev_ctx_);
            cout << "Failed to register memory" << endl;
            return -1;
        }
        mm_.set_mr(mr);
    }

    return 0;
}

struct ibv_qp* Rdma::RdmaCreateQueuepair()
{
    struct ibv_qp* qp;
    struct ibv_qp_init_attr qp_attr;
    memset(&qp_attr, 0, sizeof(qp_attr));

    /* qp type select RC for write/read */
    qp_attr.qp_type = IBV_QPT_RC;

    qp_attr.send_cq = send_cq_;
    qp_attr.recv_cq = recv_cq_;
    //if (srq_) {
    //    qp_attr.srq = srq_;
    //}

    qp_attr.cap.max_send_wr = QP_MAX_DEPTH;
    qp_attr.cap.max_recv_wr = QP_MAX_DEPTH;

    qp_attr.cap.max_send_sge = WR_MAX_SGE;
    qp_attr.cap.max_recv_sge = WR_MAX_SGE;

    qp_attr.cap.max_inline_data = MAX_INLINE_SIZE;

    qp_attr.sq_sig_all = 0;
    qp = ibv_create_qp(pd_, &qp_attr);
    if (!qp) {
        cout << "Failed to create qp " << errno << endl;
        return (struct ibv_qp*)NULL;
    }

    return qp;
}

int Rdma::RdmaDestroyQueuepair(struct ibv_qp* qp)
{
    if(ibv_destroy_qp(qp)){
        return -1;
    }
    return 0;
}

int Rdma::ModifyStateERR(struct ibv_qp* qp)
{
    if (!qp) {
        cout << "qp is NULL" << endl;
        return -1;
    }
    struct ibv_qp_attr qp_attr;
    int attr_mask;

    memset(&qp_attr, 0, sizeof(qp_attr));
    qp_attr.qp_state = IBV_QPS_ERR;
   

    attr_mask = IBV_QP_STATE;
    if (ibv_modify_qp(qp, &qp_attr, attr_mask)) {
        cout << "QP state failed to transition from * to ERR" << endl;
        return -1;
    }
    return 0;
}

int Rdma::ModifyStateINIT(struct ibv_qp* qp)
{
    if (!qp) {
        cout << "qp is NULL" << endl;
        return -1;
    }
    struct ibv_qp_attr qp_attr;
    int attr_mask;

    memset(&qp_attr, 0, sizeof(qp_attr));
    qp_attr.qp_state = IBV_QPS_INIT;
    qp_attr.port_num = resolve_port_id_;
    qp_attr.qp_access_flags = IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_ATOMIC;
    qp_attr.pkey_index = 0;

    attr_mask = IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS;
    if (ibv_modify_qp(qp, &qp_attr, attr_mask)) {
        cout << "QP state failed to transition from RESET to INIT" << endl;
        return -1;
    }
    return 0;
}

int Rdma::ModifyStateRTR(struct ibv_qp* qp, uint32_t qpn, uint16_t lid, uint8_t* gid)
{
    if (!qp) {
        cout << "qp is NULL" << endl;
        return -1;
    }
    struct ibv_qp_attr qp_attr;
    int attr_mask;

    memset(&qp_attr, 0, sizeof(qp_attr));
    qp_attr.qp_state = IBV_QPS_RTR;
    qp_attr.path_mtu = IBV_MTU_1024;
    qp_attr.dest_qp_num = qpn;
    /* set packet sequence number of recv queue as constant */
    qp_attr.rq_psn = 1234;

    qp_attr.ah_attr.dlid = lid;
    qp_attr.ah_attr.sl = 0;
    qp_attr.ah_attr.src_path_bits = 0;
    qp_attr.ah_attr.port_num = resolve_port_id_;

    /* When using RoCE, GRH must be configured */
    qp_attr.ah_attr.is_global = 1;
    memcpy(&qp_attr.ah_attr.grh.dgid, gid, 16);
    qp_attr.ah_attr.grh.flow_label = 0;
    qp_attr.ah_attr.grh.hop_limit = 1;
    qp_attr.ah_attr.grh.sgid_index = DEFAULT_GID_INDEX;
    qp_attr.ah_attr.grh.traffic_class = 0;

    qp_attr.max_dest_rd_atomic = 16;
    qp_attr.min_rnr_timer = 4;

    attr_mask = IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN
        | IBV_QP_RQ_PSN | IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER;

    if (ibv_modify_qp(qp, &qp_attr, attr_mask)) {
        cout << "QP state failed to transition from INIT to RTR" << endl;
        return -1;
    }
    return 0;
}

int Rdma::ModifyStateRTS(struct ibv_qp* qp)
{
    struct ibv_qp_attr qp_attr;
    int attr_mask;
    memset(&qp_attr, 0, sizeof(struct ibv_qp_attr));

    qp_attr.qp_state = IBV_QPS_RTS;
    qp_attr.sq_psn = 1234;
    /* rc retransmission setting */
    qp_attr.timeout = 8;
    qp_attr.retry_cnt = 7;
    qp_attr.rnr_retry = 7;//infinite
    /* atomic operation setting */
    qp_attr.max_rd_atomic = 16;
    qp_attr.max_dest_rd_atomic = 16;

    attr_mask = IBV_QP_STATE | IBV_QP_SQ_PSN | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT
        | IBV_QP_RNR_RETRY | IBV_QP_MAX_QP_RD_ATOMIC;

    if (ibv_modify_qp(qp, &qp_attr, attr_mask)) {
        cout << "QP state failed to transition from RTR to RTS" << endl;
        return -1;
    }
    return 0;
}

static inline void _print_status(struct ibv_qp* qp, const char* status)
{
    cout << "Client " << qp->qp_num << " with QP state: " << status << endl;
}

void Rdma::RdmaQueryQueuepair(struct ibv_qp* qp) {
    struct ibv_qp_attr attr;
    struct ibv_qp_init_attr init_attr;
    ibv_query_qp(qp, &attr, IBV_QP_STATE, &init_attr);
    switch (attr.qp_state) {
    case IBV_QPS_RESET:
        _print_status(qp, "IBV_QPS_RESET");
        break;
    case IBV_QPS_INIT:
        _print_status(qp, "IBV_QPS_INIT");
        break;
    case IBV_QPS_RTR:
        _print_status(qp, "IBV_QPS_RTR");
        break;
    case IBV_QPS_RTS:
        _print_status(qp, "IBV_QPS_RTS");
        break;
    case IBV_QPS_SQD:
        _print_status(qp, "IBV_QPS_SQD");
        break;
    case IBV_QPS_SQE:
        _print_status(qp, "IBV_QPS_SQE");
        break;
    case IBV_QPS_ERR:
        _print_status(qp, "IBV_QPS_ERR");
        break;
    case IBV_QPS_UNKNOWN:
        _print_status(qp, "IBV_QPS_UNKNOWN");
        break;
    }
}


int Rdma::RdmaWrite(struct ibv_qp* qp, uint32_t rkey, uint64_t source_buffer, uint64_t dest_buffer, uint64_t buffer_size, uint64_t ctx)
{
    struct ibv_sge sge;
    struct ibv_send_wr wr;
    struct ibv_send_wr* bad_wr;

    memset(&sge, 0, sizeof(struct ibv_sge));
    sge.addr = (uintptr_t)source_buffer;
    sge.length = buffer_size;
    sge.lkey = static_cast<struct ibv_mr*>(mm_.mr())->lkey;

    memset(&wr, 0, sizeof(struct ibv_send_wr));
    wr.wr_id = ctx;
    wr.sg_list = &sge;
    wr.num_sge = 1;

    wr.opcode = IBV_WR_RDMA_WRITE;
    wr.send_flags = IBV_SEND_SIGNALED;

    wr.wr.rdma.remote_addr = dest_buffer;
    wr.wr.rdma.rkey = rkey;
    if (ibv_post_send(qp, &wr, &bad_wr)) {
        cout << "Failed to post RDMA write" << endl;
        return -1;
    }
    return 0;
}

int Rdma::RdmaRead(struct ibv_qp* qp, uint32_t rkey, uint64_t source_buffer, uint64_t dest_buffer, uint64_t buffer_size, uint64_t ctx, bool signal)
{
    struct ibv_sge sge;
    struct ibv_send_wr wr;
    struct ibv_send_wr* bad_wr;

    memset(&sge, 0, sizeof(struct ibv_sge));
    sge.addr = (uintptr_t)source_buffer;
    sge.length = buffer_size;
    sge.lkey = static_cast<struct ibv_mr*>(mm_.mr())->lkey;

    memset(&wr, 0, sizeof(struct ibv_send_wr));
    wr.wr_id = ctx;
    wr.sg_list = &sge;
    wr.num_sge = 1;

    wr.opcode = IBV_WR_RDMA_READ;
    wr.send_flags = signal ? IBV_SEND_SIGNALED : 0;

    wr.wr.rdma.remote_addr = dest_buffer;
    wr.wr.rdma.rkey = rkey;
    if (ibv_post_send(qp, &wr, &bad_wr)) {
        cout << "Failed to post RDMA write" << endl;
        return -1;
    }
    return 0;

}

int Rdma::RdmaSend(struct ibv_qp* qp, uint64_t buffer, uint64_t buffer_size, uint64_t ctx)
{
    struct ibv_sge sge;
    struct ibv_send_wr wr;
    struct ibv_send_wr* bad_wr;


    memset(&sge, 0, sizeof(struct ibv_sge));
    sge.addr = (uintptr_t)buffer;
    sge.length = buffer_size;
    sge.lkey = static_cast<struct ibv_mr*>(mm_.mr())->lkey;

    memset(&wr, 0, sizeof(struct ibv_send_wr));
    /*default set sge addr*/
    wr.wr_id = ctx;

    /*only support one sge per request*/
    wr.sg_list = &sge;
    wr.num_sge = 1;


    wr.opcode = IBV_WR_SEND;
    wr.send_flags = IBV_SEND_SIGNALED;

    if (ibv_post_send(qp, &wr, &bad_wr)) {
        cout << "Failed to RDMA send" << endl;
        return 0;
    }
    return 1;
}



